"""
BLAST序列搜索工具
调用后端API /api/search/sequence/ 进行BLAST序列搜索
"""

import json
import requests
from qwen_agent.tools import BaseTool
from typing import Dict, Any
import re
import config


class BlastSequenceSearchTool(BaseTool):
    """BLAST序列搜索工具类"""
    
    # 必须设置的类属性
    name = "blast_sequence_search"
    description = "使用BLAST算法搜索与给定氨基酸序列相似的蛋白质。支持FASTA格式或纯序列输入。"
    parameters = {
        "type": "object",
        "properties": {
            "sequence": {
                "type": "string",
                "description": "氨基酸序列（支持FASTA格式或纯序列），长度10-10000个氨基酸"
            },
            "evalue": {
                "type": "number",
                "description": "E-value阈值，越小越严格（默认0.001，范围1e-10到10）",
                "default": 0.001
            },
            "max_target_seqs": {
                "type": "integer",
                "description": "最大返回结果数（默认100，范围1-500）",
                "default": 100
            }
        },
        "required": ["sequence"]
    }
    
    def _clean_sequence(self, sequence: str) -> str:
        """
        清理序列：移除FASTA头部、空白字符等
        参考后端API的clean_sequence函数逻辑
        """
        # 移除FASTA格式的头部（>开头的行）
        lines = sequence.strip().split('\n')
        clean_lines = []
        for line in lines:
            line = line.strip()
            if line and not line.startswith('>'):
                clean_lines.append(line)
        
        # 合并所有序列行
        clean_seq = ''.join(clean_lines)
        
        # 移除空白字符
        clean_seq = re.sub(r'\s+', '', clean_seq)
        
        # 验证氨基酸字符（20种标准氨基酸）
        valid_chars = set('ACDEFGHIKLMNPQRSTVWY')
        invalid_chars = set(clean_seq.upper()) - valid_chars
        
        if invalid_chars:
            raise ValueError(f"序列包含无效字符: {', '.join(invalid_chars)}")
        
        return clean_seq.upper()
    
    def call(self, params: Dict[str, Any], **kwargs) -> str:
        """
        执行BLAST序列搜索
        
        Args:
            params: 参数字典，包含sequence、evalue、max_target_seqs
            **kwargs: 其他关键字参数
            
        Returns:
            str: JSON格式的字符串结果
        """
        try:
            # 参数解析和验证
            if isinstance(params, str):
                params = json.loads(params)
            
            sequence = params.get("sequence", "").strip()
            if not sequence:
                return json.dumps({
                    "success": False,
                    "error": "序列不能为空"
                }, ensure_ascii=False)
            
            # 清理和验证序列
            try:
                clean_seq = self._clean_sequence(sequence)
            except ValueError as e:
                return json.dumps({
                    "success": False,
                    "error": str(e)
                }, ensure_ascii=False)
            
            # 验证序列长度
            if len(clean_seq) < config.LIMITS["blast"]["min_sequence_length"]:
                return json.dumps({
                    "success": False,
                    "error": f"序列太短，至少需要{config.LIMITS['blast']['min_sequence_length']}个氨基酸"
                }, ensure_ascii=False)
            
            if len(clean_seq) > config.LIMITS["blast"]["max_sequence_length"]:
                return json.dumps({
                    "success": False,
                    "error": f"序列太长，最多允许{config.LIMITS['blast']['max_sequence_length']}个氨基酸"
                }, ensure_ascii=False)
            
            # 获取其他参数
            evalue = params.get("evalue", config.DEFAULT_PARAMS["blast"]["evalue"])
            max_target_seqs = params.get("max_target_seqs", config.DEFAULT_PARAMS["blast"]["max_target_seqs"])
            
            # 参数验证
            if not (config.LIMITS["blast"]["min_evalue"] <= evalue <= config.LIMITS["blast"]["max_evalue"]):
                return json.dumps({
                    "success": False,
                    "error": f"E-value必须在{config.LIMITS['blast']['min_evalue']}到{config.LIMITS['blast']['max_evalue']}之间"
                }, ensure_ascii=False)
            
            if not (config.LIMITS["blast"]["min_max_target_seqs"] <= max_target_seqs <= config.LIMITS["blast"]["max_max_target_seqs"]):
                return json.dumps({
                    "success": False,
                    "error": f"max_target_seqs必须在{config.LIMITS['blast']['min_max_target_seqs']}到{config.LIMITS['blast']['max_max_target_seqs']}之间"
                }, ensure_ascii=False)
            
            # 构建API请求（使用Form Data）
            api_data = {
                "sequence": clean_seq,
                "evalue": evalue,
                "max_target_seqs": max_target_seqs
            }
            
            # 调用API（POST请求，超时时间包含BLAST执行时间）
            response = requests.post(
                config.API_ENDPOINTS["blast"],
                data=api_data,
                timeout=config.LIMITS["blast"]["timeout"]
            )
            response.raise_for_status()
            api_result = response.json()
            
            # 格式化结果
            formatted_result = {
                "success": api_result.get("success", True),
                "message": api_result.get("message", ""),
                "query_length": api_result.get("query_length", len(clean_seq)),
                "total_hits": api_result.get("total_hits", 0),
                "execution_time": api_result.get("execution_time", 0),
                "cached": api_result.get("cached", False),
                "results": []
            }
            
            # 格式化结果列表
            for item in api_result.get("results", []):
                formatted_result["results"].append({
                    "id": item.get("id"),
                    "accession": item.get("accession"),
                    "species": item.get("species"),
                    "type": item.get("type"),
                    "assembly": item.get("assembly"),
                    "has_pdb": item.get("has_pdb", False),
                    "pident": item.get("identity", item.get("pident")),  # 一致性百分比
                    "length": item.get("alignment_length", item.get("length")),
                    "mismatch": item.get("mismatches", item.get("mismatch")),
                    "gapopen": item.get("gap_opens", item.get("gapopen")),
                    "evalue": item.get("evalue"),
                    "bitscore": item.get("bit_score", item.get("bitscore")),
                    "qstart": item.get("query_start", item.get("qstart")),
                    "qend": item.get("query_end", item.get("qend")),
                    "sstart": item.get("subject_start", item.get("sstart")),
                    "send": item.get("subject_end", item.get("send"))
                })
            
            return json.dumps(formatted_result, ensure_ascii=False)
            
        except requests.exceptions.Timeout:
            error_msg = {
                "success": False,
                "error": f"搜索超时（>{config.LIMITS['blast']['timeout']}秒），请缩短序列或增大E-value"
            }
            return json.dumps(error_msg, ensure_ascii=False)
        
        except requests.exceptions.RequestException as e:
            error_msg = {
                "success": False,
                "error": f"API调用失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)
        
        except Exception as e:
            error_msg = {
                "success": False,
                "error": f"工具执行失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)

